{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7d96fb89",
   "metadata": {},
   "source": [
    "[参考链接-唐国梁Tommy-轻松学 PyTorch 手写字体识别 MNIST](https://www.bilibili.com/video/BV1WT4y177SA)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2af90bb",
   "metadata": {},
   "source": [
    "### 1. 加载必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "81f6c144",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df02de9e",
   "metadata": {},
   "source": [
    "### 2. 定义超参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "14e8ad4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 16   # 每批处理的数据\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 是否用GPU\n",
    "EPOCHS = 10   # 训练数据集的轮次"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "905064aa",
   "metadata": {},
   "source": [
    "### 3. 构建`pipeline`，对图像做处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4cc1dd01",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = transforms.Compose([\n",
    "    transforms.ToTensor(),   #  将图片转换为 Tensor\n",
    "    transforms.Normalize((0.1307,), (0.3081,))  # 正则化，降低模型复杂度，参数值官网查询所得，直接使用即可\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d09cc3f3",
   "metadata": {},
   "source": [
    "### 4. 下载，加载数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2375e652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# 下载数据集\n",
    "train_set = datasets.MNIST('./data', train=True, download=True, transform=pipeline)\n",
    "\n",
    "test_set = datasets.MNIST('./data', train=False, download=True, transform=pipeline)\n",
    "\n",
    "# 加载数据\n",
    "train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)\n",
    "\n",
    "test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5b4ebfc",
   "metadata": {},
   "source": [
    "#### 插入代码，显示MNIST中的图片"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7a04e4c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/MNIST/raw/train-images-idx3-ubyte', 'rb') as f:\n",
    "    file = f.read()\n",
    "    \n",
    "image1 = [int(str(item).encode('ascii'), 16) for item in file[16 : 16+784]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6aa856e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(28, 28, 1)\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "image1_np = np.array(image1, dtype=np.uint8).reshape(28, 28, 1) # 1: 通道，灰度图片\n",
    "\n",
    "print(image1_np.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "16ee5fc8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图片保存\n",
    "cv2.imwrite('digit.jpg', image1_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "317b6ed6",
   "metadata": {},
   "source": [
    "### 5. 构建网络模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d2cd2cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Digit(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, 5)  # 1：灰度图的通道， 10：输出通道， 5： kernel size\n",
    "        self.conv2 = nn.Conv2d(10, 20, 3)  # 10： 输入通道， 20：输出通道， 3：kernel size\n",
    "        self.fc1 = nn.Linear(20*10*10, 500)  # 20*10*10： 输入通道， 500：输出通道\n",
    "        self.fc2 = nn.Linear(500, 10)   # 500：输入通道， 10：输出通道（10种数字0~9）\n",
    "        \n",
    "    def forward(self, x):\n",
    "        input_size = x.size(0)  # 0 : batch_size,   [batch_size*1*28*28]\n",
    "        x = self.conv1(x)  # 输入： batch*1*28*28, 输出： batch*10*24*24 （28-5+1=24）\n",
    "        x = F.relu(x)  # 保持shape 不变，输出：bacth*10*24*24\n",
    "        x = F.max_pool2d(x, 2, 2)  # 输入： bacth*10*24*24  输出： bacth*10*12*12\n",
    "        \n",
    "        x = self.conv2(x)  # 输入： bacth*10*12*12  输出： batch*20*10*10 （12-3+1=10）\n",
    "        x = F.relu(x)  \n",
    "        \n",
    "        x = x.view(input_size, -1)  # 拉平， -1 自动计算维度， 20*10*10 = 2000\n",
    "        \n",
    "        x = self.fc1(x)  # 输入：batch*2000   输出：batch*500\n",
    "        x = F.relu(x)  # 保持 shape 不变\n",
    "        \n",
    "        x = self.fc2(x)  # 输入： batch*500   输出： batch*10\n",
    "        \n",
    "        output = F.log_softmax(x, dim=1)  # 计算分类后，每个数字的概率值\n",
    "        \n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81eeb5b",
   "metadata": {},
   "source": [
    "### 6. 定义优化器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1039220c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Digit().to(DEVICE)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7fca7aa",
   "metadata": {},
   "source": [
    "### 7. 定义训练方法"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e391115f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, device, train_loader, optimizer, epoch):\n",
    "    # 模型训练\n",
    "    model.train()\n",
    "    for batch_index, (data, target) in enumerate(train_loader):\n",
    "        # 部署到 DEVICE 上去\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        # 梯度初始化为 0\n",
    "        optimizer.zero_grad()\n",
    "        # 训练后的结果\n",
    "        output = model(data)\n",
    "        # 计算损失\n",
    "        loss = F.cross_entropy(output, target)\n",
    "        # 反向传播\n",
    "        loss.backward()\n",
    "        # 参数优化\n",
    "        optimizer.step()\n",
    "        \n",
    "        if batch_index % 3000 == 0:\n",
    "            print('Train epoch : {} \\t loss: {:.6f}'.format(epoch, loss.item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "931127d4",
   "metadata": {},
   "source": [
    "### 8. 定义测试方法"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6f2a26fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, device, test_loader):\n",
    "    # 模型验证\n",
    "    model.eval()\n",
    "    # 正确率\n",
    "    correct = 0.0\n",
    "    # 测试损失\n",
    "    test_loss = 0.0\n",
    "    with torch.no_grad():  # 不会计算梯度，也不会进行反向传播\n",
    "        for data, target in test_loader:\n",
    "            # 部署到 DEVICE 上\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            # 测试数据\n",
    "            output = model(data)\n",
    "            # 计算测试损失\n",
    "            test_loss += F.cross_entropy(output, target).item() \n",
    "            # 找到概率值最大的下标，三种写法一样\n",
    "            pred = output.max(1, keepdim=True)[1]  # 值， 索引\n",
    "            # pred = torch.max(output, dim=1)  # 1: 横轴\n",
    "            # pred = output.argmax(dim=1)\n",
    "    \n",
    "            # 累计正确的值\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "        test_loss /= len(test_loader.dataset)\n",
    "        \n",
    "        print('Test -- Average loss: {:.4f}, accuracy: {:.3f}\\n'.format(\n",
    "            test_loss, 100.0*correct/ len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce1cb7c0",
   "metadata": {},
   "source": [
    "### 9. 方法调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e5e8b4f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train epoch : 1 \t loss: 2.310601\n",
      "Train epoch : 1 \t loss: 0.002911\n",
      "Test -- Average loss: 0.0027, accuracy: 98.530\n",
      "\n",
      "Train epoch : 2 \t loss: 0.000667\n",
      "Train epoch : 2 \t loss: 0.015995\n",
      "Test -- Average loss: 0.0025, accuracy: 98.740\n",
      "\n",
      "Train epoch : 3 \t loss: 0.001822\n",
      "Train epoch : 3 \t loss: 0.000057\n",
      "Test -- Average loss: 0.0024, accuracy: 98.850\n",
      "\n",
      "Train epoch : 4 \t loss: 0.003403\n",
      "Train epoch : 4 \t loss: 0.002398\n",
      "Test -- Average loss: 0.0032, accuracy: 98.610\n",
      "\n",
      "Train epoch : 5 \t loss: 0.005771\n",
      "Train epoch : 5 \t loss: 0.000003\n",
      "Test -- Average loss: 0.0034, accuracy: 98.730\n",
      "\n",
      "Train epoch : 6 \t loss: 0.000492\n",
      "Train epoch : 6 \t loss: 0.052056\n",
      "Test -- Average loss: 0.0027, accuracy: 98.920\n",
      "\n",
      "Train epoch : 7 \t loss: 0.001504\n",
      "Train epoch : 7 \t loss: 0.000004\n",
      "Test -- Average loss: 0.0030, accuracy: 98.830\n",
      "\n",
      "Train epoch : 8 \t loss: 0.000100\n",
      "Train epoch : 8 \t loss: 0.000110\n",
      "Test -- Average loss: 0.0036, accuracy: 98.880\n",
      "\n",
      "Train epoch : 9 \t loss: 0.000687\n",
      "Train epoch : 9 \t loss: 0.000079\n",
      "Test -- Average loss: 0.0046, accuracy: 98.800\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(1, EPOCHS):\n",
    "    train_model(model, DEVICE, train_loader, optimizer, epoch)\n",
    "    test_model(model, DEVICE, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f0ac487",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-cv",
   "language": "python",
   "name": "pytorch-cv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
